import os
import sys
import onnx
from onnx import helper
import onnx_graphsurgeon as gs
import numpy as np
from typing import List, Set, Dict, Tuple


# === 工程路径配置 ===
PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
if PROJECT_ROOT not in sys.path:
    sys.path.insert(0, PROJECT_ROOT)

# 路径配置
NMS_ONNX = os.path.join(PROJECT_ROOT, 'onnx/yolov11n_nms.onnx')
MERGE_INSERT_NODE = "/Transpose"  # 分割节点名


a = os.path.join(PROJECT_ROOT, 'onnx/nms_before.onnx')
b = os.path.join(PROJECT_ROOT, 'onnx/yolov11n_nms.onnx')


def detailed_model_analysis(model1_path, model2_path):
    """
    详细分析两个模型的结构差异
    """
    model1 = onnx.load(model1_path)
    model2 = onnx.load(model2_path)
    
    print("模型结构分析:")
    print(f"模型1节点数: {len(model1.graph.node)}")
    print(f"模型2节点数: {len(model2.graph.node)}")
    
    # 比较节点类型序列
    node_types1 = [node.op_type for node in model1.graph.node]
    node_types2 = [node.op_type for node in model2.graph.node]
    
    flag = False
    
    if node_types1 == node_types2:
        print("✓ 节点类型序列相同")
    else:
        print("✗ 节点类型序列不同")
        for i, (t1, t2) in enumerate(zip(node_types1, node_types2)):
            if t1 != t2:
                if(flag == False):
                    flag = True
                # print(f"  节点 {i}: {t1} vs {t2}")
                
            if(flag):
                print(f"  节点 {i}: {t1} vs {t2}")
    
    # 比较参数数量
    params1 = len(model1.graph.initializer)
    params2 = len(model2.graph.initializer)
    print(f"参数数量: {params1} vs {params2}")

# 运行分析
detailed_model_analysis(a, b)